package top.atcra.rpc.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import top.atcra.rpc.config.ConfigLoader;
import top.atcra.rpc.exception.ZookeeperException;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;

@Slf4j
public class ZookeeperUtil {
    private static final Map<String, ZooKeeper> zkMap = new ConcurrentHashMap<>();

    public static ZooKeeper createZookeeper(String connectStr, int timeout_ms) {
        CountDownLatch countDownLatch = new CountDownLatch(1);
        try {
            return new ZooKeeper(connectStr, timeout_ms, null);
        } catch (IOException e) {
            throw new ZookeeperException(e);
        } finally {
            countDownLatch.countDown();
        }
    }

    public static ZooKeeper createZookeeper(String connectStr) {
        if (zkMap.containsKey(connectStr)) {
            return zkMap.get(connectStr);
        }
        var cfg = ConfigLoader.getConfig().getZookeeper();
        var zk = createZookeeper(connectStr, cfg.getTimeout_ms());
        zkMap.put(connectStr, zk);
        return zk;
    }


    /**
     * 关闭zookeeper的方法
     *
     * @param zooKeeper zooKeeper实例
     */
    public static void close(ZooKeeper zooKeeper) {
        try {
            zooKeeper.close();
        } catch (InterruptedException e) {
            log.error("关闭zookeeper时发生问题：", e);
            throw new ZookeeperException(e);
        }
    }

    /**
     * 查询一个节点的子元素
     *
     * @param zooKeeper   zk实例
     * @param serviceNode 服务节点
     * @return 子元素列表
     */
    public static List<String> getChildren(ZooKeeper zooKeeper, String serviceNode, Watcher watcher) {
        try {
            return zooKeeper.getChildren(serviceNode, watcher);
        } catch (KeeperException | InterruptedException e) {
            log.error("获取节点【{}】的子元素时发生异常.", serviceNode, e);
            throw new ZookeeperException(e);
        }
    }

    /**
     * 判断节点是否存在
     *
     * @param zk      zk实例
     * @param node    节点路劲
     * @param watcher watcher
     * @return ture 存在 | false 不存在
     */
    public static boolean exists(ZooKeeper zk, String node, Watcher watcher) {
        try {
            return zk.exists(node, watcher) != null;
        } catch (KeeperException | InterruptedException e) {
            log.error("判断节点[{}]是否存在是发生异常", node, e);
            throw new ZookeeperException(e);
        }
    }

    /**
     * 创建一个节点的工具方法
     *
     * @param zooKeeper  zooKeeper实例
     * @param nodePath   节点路径
     * @param data       节点存储的数据
     * @param watcher    watcher实例
     * @param createMode 节点的类型
     * @return true: 成功创建  false: 已经存在  异常：抛出
     */
    public static Boolean createNode(ZooKeeper zooKeeper, String nodePath, byte[] data, Watcher watcher, CreateMode createMode) {
        try {
            if (zooKeeper.exists(nodePath, watcher) == null) {
                String result = zooKeeper.create(nodePath, data,
                        ZooDefs.Ids.OPEN_ACL_UNSAFE, createMode);
                log.info("节点【{}】，成功创建。", result);
                return true;
            } else {
                if (log.isDebugEnabled()) {
                    log.info("节点【{}】已经存在，无需创建。", nodePath);
                }
                return false;
            }
        } catch (KeeperException | InterruptedException e) {
            log.error("创建目录时发生异常：", e);
            throw new ZookeeperException(e);
        }
    }

    /**
     * 若结点不存在，创建一个节点的工具方法
     *
     * @param zooKeeper  zooKeeper实例
     * @param nodePath   节点路径
     * @param data       节点存储的数据
     * @param watcher    watcher实例
     * @param createMode 节点的类型
     * @return true: 成功创建  false: 已经存在  异常：抛出
     */
    public static Boolean createNodeIfNotExist(ZooKeeper zooKeeper, String nodePath, byte[] data, Watcher watcher, CreateMode createMode) {
        if (!exists(zooKeeper, nodePath, watcher)) {
            return createNode(zooKeeper, nodePath, data, watcher, createMode);
        }
        return false;
    }

    /**
     * 创建一个节点的工具方法(data=null)
     *
     * @param zooKeeper  zooKeeper实例
     * @param nodePath   节点路径
     * @param watcher    watcher实例
     * @param createMode 节点的类型
     * @return true: 成功创建  false: 已经存在  异常：抛出
     */
    public static Boolean createNode(ZooKeeper zooKeeper, String nodePath, Watcher watcher, CreateMode createMode) {
        return createNode(zooKeeper, nodePath, null, watcher, createMode);
    }

    /**
     * 若结点不存在，创建一个节点的工具方法(data=null)
     *
     * @param zooKeeper  zooKeeper实例
     * @param nodePath   节点路径
     * @param watcher    watcher实例
     * @param createMode 节点的类型
     * @return true: 成功创建  false: 已经存在  异常：抛出
     */
    public static Boolean createNodeIfNotExist(ZooKeeper zooKeeper, String nodePath, Watcher watcher, CreateMode createMode) {
        if (!exists(zooKeeper, nodePath, watcher)) {
            return createNode(zooKeeper, nodePath, null, watcher, createMode);
        }
        return false;
    }

}
